import torch
import torch.nn.functional as F
import torch.nn as nn
import math
from torch.nn import MaxPool2d


class ModuleWithAttr(nn.Module):

    # 只能是数字，默认注册为0

    def __init__(self, extra_info=['step']):
        super(ModuleWithAttr, self).__init__()
        for key in extra_info:
            self.set_buffer(key, 0)

    def set_buffer(self, key, value):
        if not(hasattr(self, '__' + key)):
            self.register_buffer('__' + key, torch.tensor(value))
        setattr(self, '__' + key, torch.tensor(value))

    def get_buffer(self, key):
        if not(hasattr(self, '__' + key)):
            raise Exception('no such key!')
        return getattr(self, '__' + key).item()


def find_idx(list_, name):
    for i in range(len(list_)):
        if list_[i] == name:
            return i


class MCNN(ModuleWithAttr):

    def __init__(self, label_name):
        super(MCNN, self).__init__()
        self.label_name = label_name
        self.groups = [
            ['Male', ],
            ['Big_Nose', 'Pointy_Nose', ],
            ['Big_Lips', 'Smiling', 'Wearing_Lipstick', 'Mouth_Slightly_Open', ],
            ['Arched_Eyebrows', 'Bags_Under_Eyes', 'Bushy_Eyebrows', 'Narrow_Eyes', 'Eyeglasses'],
            ['Attractive', 'Blurry', 'Oval_Face', 'Pale_Skin', 'Young', 'Heavy_Makeup', ],
            ['Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Gray_Hair', 'Wearing_Earrings', 'Wearing_Necklace',
             'Wearing_Necktie', 'Bald', 'Receding_Hairline', 'Bangs', 'Wearing_Hat', 'Straight_Hair', 'Wavy_Hair'],
            ['5_o_Clock_Shadow', 'Mustache', 'No_Beard', 'Sideburns', 'Goatee'],
            ['High_Cheekbones', 'Rosy_Cheeks', ],
            ['Chubby', 'Double_Chin'],
        ]
        self.num_group = len(self.groups)

        self.g2num = [[find_idx(self.label_name, name) for name in group] for group in self.groups]

        self.shared = nn.Sequential(
            nn.Conv2d(3, 75, 7, stride=1, padding=0),
            nn.ReLU(),
            nn.MaxPool2d(3),
            nn.Conv2d(75, 200, 5, stride=1, padding=0),
            nn.ReLU(),
            nn.MaxPool2d(3),
        )

        conv3s = []
        fcs = []
        for i in range(self.num_group):
            conv3s.append(nn.Sequential(
                nn.Conv2d(200, 300, 3),
                nn.ReLU(),
                nn.MaxPool2d(5),
            ))
            fcs.append(nn.Sequential(
                nn.Linear(4800, 512),
                nn.ReLU(),
                nn.Dropout(),
                nn.Linear(512, 512),
                nn.ReLU(),
                nn.Dropout(),
                nn.Linear(512, len(self.groups[i]) * 2)
            ))
        self.conv3s = nn.ModuleList(conv3s)
        self.fcs = nn.ModuleList(fcs)

        self.label_loss = nn.CrossEntropyLoss()

    def forward(self, sample):
        x = self.shared(sample['image'])
        output = [None] * len(self.label_name)
        for i in range(self.num_group):
            group_output = self.conv3s[i](x)
            group_output = group_output.reshape(group_output.size()[0], -1)
            group_output = self.fcs[i](group_output)
            group_output = group_output.reshape(group_output.size()[0], 2, -1)
            group_output = torch.split(group_output, 1, 2)
            # 把group的内容
            for j, value in enumerate(group_output):
                output[self.g2num[i][j]] = value
        output = torch.cat(output, dim=2)
        return output

    def cal_loss(self, sample, output):
        loss = {
            'all': self.label_loss(output, sample['label'])
        }
        return loss


class MCNN_BN(ModuleWithAttr):

    def __init__(self, label_name):
        super(MCNN_BN, self).__init__()
        self.label_name = label_name
        self.groups = [
            ['Male', ],
            ['Big_Nose', 'Pointy_Nose', ],
            ['Big_Lips', 'Smiling', 'Wearing_Lipstick', 'Mouth_Slightly_Open', ],
            ['Arched_Eyebrows', 'Bags_Under_Eyes', 'Bushy_Eyebrows', 'Narrow_Eyes', 'Eyeglasses'],
            ['Attractive', 'Blurry', 'Oval_Face', 'Pale_Skin', 'Young', 'Heavy_Makeup', ],
            ['Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Gray_Hair', 'Wearing_Earrings', 'Wearing_Necklace',
             'Wearing_Necktie', 'Bald', 'Receding_Hairline', 'Bangs', 'Wearing_Hat', 'Straight_Hair', 'Wavy_Hair'],
            ['5_o_Clock_Shadow', 'Mustache', 'No_Beard', 'Sideburns', 'Goatee'],
            ['High_Cheekbones', 'Rosy_Cheeks', ],
            ['Chubby', 'Double_Chin'],
        ]
        self.num_group = len(self.groups)

        self.g2num = [[find_idx(self.label_name, name) for name in group] for group in self.groups]

        self.shared = nn.Sequential(
            nn.Conv2d(3, 75, 7, stride=1, padding=0),
            nn.BatchNorm2d(75),
            nn.ReLU(),
            nn.MaxPool2d(3),
            nn.Conv2d(75, 200, 5, stride=1, padding=0),
            nn.BatchNorm2d(200),
            nn.ReLU(),
            nn.MaxPool2d(3),
        )

        conv3s = []
        fcs = []
        for i in range(self.num_group):
            conv3s.append(nn.Sequential(
                nn.Conv2d(200, 300, 3),
                nn.BatchNorm2d(300),
                nn.ReLU(),
                nn.MaxPool2d(5),
            ))
            fcs.append(nn.Sequential(
                nn.Linear(4800, 512),
                nn.ReLU(),
                nn.Dropout(),
                nn.Linear(512, 512),
                nn.ReLU(),
                nn.Dropout(),
                nn.Linear(512, len(self.groups[i]) * 2)
            ))
        self.conv3s = nn.ModuleList(conv3s)
        self.fcs = nn.ModuleList(fcs)

        self.label_loss = nn.CrossEntropyLoss()

    def forward(self, sample):
        x = self.shared(sample['image'])
        output = [None] * len(self.label_name)
        for i in range(self.num_group):
            group_output = self.conv3s[i](x)
            group_output = group_output.reshape(group_output.size()[0], -1)
            group_output = self.fcs[i](group_output)
            group_output = group_output.reshape(group_output.size()[0], 2, -1)
            group_output = torch.split(group_output, 1, 2)
            # 把group的内容
            for j, value in enumerate(group_output):
                output[self.g2num[i][j]] = value
        output = torch.cat(output, dim=2)
        return output

    def cal_loss(self, sample, output):
        loss = {
            'all': self.label_loss(output, sample['label'])
        }
        return loss


class Conv_BN_Pooling(nn.Module):

    def __init__(self, in_channels, out_channels, conv_kernel, conv_padding, pooling_kernel=None):
        super(Conv_BN_Pooling, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, conv_kernel, stride=1, padding=conv_padding)
        self.bn = nn.BatchNorm2d(out_channels)
        if pooling_kernel is not None:
            self.has_pooling = True
            self.pooling = nn.MaxPool2d(pooling_kernel)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        if self.has_pooling:
            x = self.pooling(x)
        return x


def func_partially_shared(SNet, TSNets, SFeature, TSFeatures):
    '''
    SNet: 单独的Module
    TSNets：TSNet的list或ModuleList(所有net，单独一层)
    SFeature：上一层SNet的输出
    TSFeatures：上一层TSNet的输出的list
    '''
    assert len(TSFeatures) == len(TSNets)
    Out_TSFeatures = []
    for i in range(len(TSFeatures)):
        Out_TSFeatures.append(
            TSNets[i](
                torch.cat([SFeature, TSFeatures[i]], dim=1)
            )
        )
    Out_SFeature = SNet(
        torch.cat([SFeature] + TSFeatures, dim=1)
    )
    return Out_SFeature, Out_TSFeatures


class MCNN_PS(ModuleWithAttr):
    def __init__(self, label_name):
        super(MCNN_PS, self).__init__()
        self.label_name = label_name

        self.C_TS = [32, 64, 128, 256, 128]
        self.C_S = [8, 16, 32, 64, 32]
        self.groups = [
            ['Arched_Eyebrows', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Bushy_Eyebrows',
             'Eyeglasses', 'Gray_Hair', 'Narrow_Eyes', 'Receding_Hairline', 'Wearing_Hat'],
            ['Big_Nose', 'High_Cheekbones', 'Pointy_Nose', 'Rosy_Cheeks', 'Sideburns', 'Wearing_Earrings'],
            ['Big_Lips', 'Double_Chin', 'Goatee', 'Mustache',
             'Mouth_Slightly_Open', 'No_Beard', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie'],
            ['5_o_Clock_Shadow', 'Attractive', 'Blurry', 'Chubby', 'Heavy_Makeup', 'Male', 'Oval_Face', 'Pale_Skin',
             'Straight_Hair', 'Smiling', 'Wavy_Hair', 'Young'],
        ]
        self.label_name = label_name
        self.g2num = [[find_idx(self.label_name, name) for name in group] for group in self.groups]

        # 第一层
        S_List = [Conv_BN_Pooling(3, self.C_S[0], 3, 1, 2)]
        # 后面的几层
        for i in range(1, len(self.C_S)):
            S_List.append(Conv_BN_Pooling(self.C_S[i - 1] + self.C_TS[i - 1] * 4, self.C_S[i], 3, 1, 2))

        TS_List = []
        for i in range(4):
            tmp = []
            tmp.append(Conv_BN_Pooling(3, self.C_TS[0], 3, 1, 2))
            for j in range(1, len(self.C_TS)):
                tmp.append(Conv_BN_Pooling(self.C_S[j - 1] + self.C_TS[j - 1], self.C_TS[j], 3, 1, 2))
            TS_List.append(nn.ModuleList(tmp))

        FC_List = []
        for i in range(4):
            FC_List.append(
                nn.Sequential(
                    nn.Linear(5 * 6 * 128, 512),
                    nn.BatchNorm1d(512),
                    nn.Linear(512, 512),
                    nn.BatchNorm1d(512),
                    nn.Linear(512, len(self.groups[i]) * 2)
                )
            )

        self.S_List = nn.ModuleList(S_List)
        self.TS_List = nn.ModuleList(TS_List)
        self.FC_List = nn.ModuleList(FC_List)

        self.label_loss = nn.CrossEntropyLoss()

    def forward(self, sample):
        # forward
        # 第一层
        Out_SFeature = self.S_List[0](sample['image'])
        Out_TSFeatures = [self.TS_List[i][0](sample['image']) for i in range(4)]
        # 后面几层
        for i in range(1, len(self.S_List)):
            Out_SFeature, Out_TSFeatures = func_partially_shared(
                self.S_List[i],
                [self.TS_List[layer][i] for layer in range(4)],
                Out_SFeature,
                Out_TSFeatures,
            )

        output = [None] * len(self.label_name)
        for i in range(4):
            group_output = Out_TSFeatures[i].reshape(Out_TSFeatures[i].shape[0], -1)
            group_output = self.FC_List[i](group_output)
            group_output = group_output.reshape(group_output.size()[0], 2, -1)
            group_output = torch.split(group_output, 1, 2)
            # 把group的内容
            for j, value in enumerate(group_output):
                output[self.g2num[i][j]] = value
        output = torch.cat(output, dim=2)
        return output

    def cal_loss(self, sample, output):
        loss = {
            'all': self.label_loss(output, sample['label'])
        }
        return loss
